# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import itertools
import json
import logging
import typing
from concurrent.futures import ThreadPoolExecutor

import absl.app
import absl.flags as flags
import tornado.ioloop

from coin.base.hostname import get_hostname
from coin.base.timestamp import get_timestamp
import coin.flow.subscriber as flow_subscriber
from coin.base.param_util import to_list
from coin.flow.simple_queue import SimpleQueue

from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.subscription import SubscriptionRequest
from coin.support.exchange_status.logic.db_importer import ExchangeStatusImporter
import coin.proto.exchange_status_pb2 as coin_checker
from xunkemgmt.support.accounting.logic.query_request import get_exchange_api_id_from_request
from coin.exchange.util.feed_checker import FeedCheckerProductConfig


def get_exchange_api_id_map(request_list):
  exchange_id_list = get_exchange_api_id_from_request(request_list)
  res = {}
  for request, exchange_id in zip(request_list, exchange_id_list):
    key = SubscriptionRequest(request.market_type, request.exchange, request.api_version)
    res[key] = exchange_id
  return res


class LiveFeedChecker(object):
  def __init__(self, feed_product, exchange_api_id_map, db_config, frequency, logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._feed_products = to_list(feed_product)
    # We use simple configuration parameters for now, which is consistent to
    # the previous version.
    feed_checker_option = FeedCheckerProductConfig(max_book_gap=30e9,
                                                   max_book_change_gap=60e9,
                                                   max_trade_gap=30e9,
                                                   max_trade_change_gap=60e9)
    feed_checker_config = {product: feed_checker_option for product in self._feed_products}
    self._feed_subsystem = MultiFeedSubsystem(feed_checker_config)
    self._exchange_api_id_map = exchange_api_id_map
    self._db_config = db_config
    self._ioloop = tornado.ioloop.IOLoop.current()
    self._queue = SimpleQueue('queue')
    self._flow_sub = flow_subscriber.from_queue(self._queue)
    self._executor = ThreadPoolExecutor(max_workers=5)
    self._frequency = frequency

  def _init_feed_subsystem(self):
    self._logger.info('Initializing feed...')
    self._feed_subsystem.subscribe_book(self._feed_products, lambda x: None)
    self._feed_subsystem.subscribe_trade(self._feed_products, lambda x: None)
    self._feed_subsystem.subscribe_feed_status(self._feed_products, self._on_feed_status)
    self._feed_subsystem.start(self._ioloop, self._queue, self._flow_sub, self._frequency)

  def _on_feed_status(self, data):
    request_list = []
    response_list = []

    exchange_request = data.pop('exchange_request')
    exchange_api_id = self._exchange_api_id_map[exchange_request]

    for product, status in data.items():
      request = coin_checker.ExchangeStatusRequestProto(type=coin_checker.FEED_REQUEST,
                                                        exchange_api_id=exchange_api_id,
                                                        symbol=product.symbol,
                                                        feed_name='websocket',
                                                        hostname=get_hostname(),
                                                        query_timestamp=get_timestamp())
      request_list.append(request)
      response_list.append(status)

    future = self._executor.submit(self.write_to_db, request_list, response_list)
    future.add_done_callback(self.future_callback)

  def write_to_db(self, request_list, response_list):
    importer = ExchangeStatusImporter.create_connector(db_config=self._db_config)
    importer.insert_feed_status_history(request_list, response_list)

  def future_callback(self, future):
    if future.exception():
      self._logger.error('Failed to write to database!')

  def start(self):
    self._ioloop.add_callback(self._init_feed_subsystem)

  def loop(self, exit_after: typing.Optional[datetime.timedelta] = None):
    if exit_after is not None:
      self._ioloop.add_timeout(exit_after, self._ioloop.stop)
    self._ioloop.start()


def main(_):
  exchange_id_request = []

  bitflyer_products = ["BTC-JPY.PERPETUAL"]  # Bitflyer can only subscribe one symbol at a time.
  bitflyer_products = [BitflyerFuturesProduct.FromStr(symbol) for symbol in bitflyer_products]
  exchange_id_request.append(
      SubscriptionRequest(
          market_type='Futures',
          exchange='Bitflyer',
          api_version=None,
      ))

  bitmex_products = ["BTC-USD.PERPETUAL"]
  bitmex_products = [BitmexFuturesProduct.FromStr(symbol) for symbol in bitmex_products]
  exchange_id_request.append(
      SubscriptionRequest(market_type='Futures', exchange='Bitmex', api_version=None))

  okexfutures_products = ["BTC-USD.THIS_WEEK"]
  okexfutures_products = [OkexFuturesProduct.FromStr(symbol) for symbol in okexfutures_products]
  exchange_id_request.append(
      SubscriptionRequest(
          market_type='Futures',
          exchange='Okex',
          api_version='futures',
      ))

  products = [
      bitflyer_products,
      bitmex_products,
      okexfutures_products,
  ]
  # Flatten product list.
  products = list(itertools.chain.from_iterable(products))

  with open(flags.FLAGS.db_config) as db_config_file:
    db_config = json.load(db_config_file)

  if flags.FLAGS.duration:
    tornado.ioloop.IOLoop.current().call_later(flags.FLAGS.duration,
                                               tornado.ioloop.IOLoop.current().stop)

  exchange_id_map = get_exchange_api_id_map(exchange_id_request)
  checker = LiveFeedChecker(products, exchange_id_map, db_config, flags.FLAGS.frequency)
  checker.start()
  checker.loop()


if __name__ == "__main__":
  flags.DEFINE_integer(
      'duration',
      None,
      'Duration for running checker in seconds. If set to None, the program will '
      'run for ever, otherwise will stop after |duration| seconds.')
  flags.DEFINE_integer('frequency', 1, 'Feed check frequency.')

  absl.app.run(main)
